#include "codec.h"
//const
char* RDIP;
int OUTWIDTH, OUTHEIGHT, FPS, RDPORT;

//global
AVFormatContext *fmtCtx;
AVStream *videoStream;
AVCodec *videoCodec;
AVCodecContext *ctx;
AVStream *st;

void set(int width, int height, int fps, char* ip, int port) {
    OUTWIDTH = width;
    OUTHEIGHT = height;
    FPS = fps;
    RDIP = ip;
    RDPORT = port;
}

AVStream* add_video_stream(AVFormatContext *oc, AVCodec **codec, enum AVCodecID codec_id) {
    //find video encoder
    *codec = avcodec_find_encoder(codec_id);
    st = avformat_new_stream(oc, *codec);
    ctx = st->codec;
    avcodec_get_context_defaults3(ctx, *codec);
    ctx->codec_id = codec_id;
    ctx->width = OUTWIDTH;
    ctx->height = OUTHEIGHT;
    ctx->time_base.den = FPS;
    ctx->time_base.num = 1;
    ctx->pix_fmt = PIX_FMT_YUV420P;
    if (oc->oformat->flags & AVFMT_GLOBALHEADER) {
        ctx->flags |= CODEC_FLAG_GLOBAL_HEADER;
    }
    av_opt_set(ctx->priv_data, "preset", "ultrafast", 0);
    av_opt_set(ctx->priv_data, "tune", "stillimage,fastdecode,zerolatency", 0);
    av_opt_set(ctx->priv_data, "x264opts", "crf=26:vbv-maxrate=728:vbv-bufsize=364:keyint=25", 0);
    return st;
}

void openNetStream() {
    snprintf(fmtCtx->filename, sizeof(fmtCtx->filename), "rtp://%s:%d", RDIP, RDPORT);
    avio_open(&fmtCtx->pb, fmtCtx->filename, AVIO_FLAG_WRITE);
    videoStream = NULL;
    videoStream = add_video_stream(fmtCtx, &videoCodec, AV_CODEC_ID_H264);
    //if the rtp address is invalid, the next line will invoke panic
    avcodec_open2(videoStream->codec, videoCodec, NULL);
    avformat_write_header(fmtCtx, NULL);
}

void fill_yuv_image(AVPicture *pict,int frame_index,int width, int height) {
    int x, y, i;
    i = frame_index;
    /* Y */
    for (y = 0; y < height; y++)
        for (x = 0; x < width; x++)
            pict->data[0][y * pict->linesize[0] +x] = x + y + i * 3;
    /* Cb and Cr */
    for (y = 0; y < height / 2; y++) {
        for (x = 0; x < width / 2; x++) {
            pict->data[1][y * pict->linesize[1] +x] = 128 + y + i * 2;
            pict->data[2][y * pict->linesize[2] +x] = 64 + x + i * 5;
        }
    }
}


void work() {
    av_register_all();
    avformat_network_init();
    //initial container
    fmtCtx = avformat_alloc_context();
    //acquire output format,then open rtp net stream
    fmtCtx->oformat = av_guess_format("rtp", NULL, NULL);
    openNetStream();
    //start send data
    AVFrame* m_pYUVFrame = av_frame_alloc();
    while(1) {
        fill_yuv_image(m_pYUVFrame, videoStream->codec->frame_number, OUTWIDTH, OUTHEIGHT);
        //encode image
        AVPacket pkt;
        int got_output = 0;
        av_init_packet(&pkt);
        pkt.data = NULL;
        pkt.size = 0;
        pkt.pts = AV_NOPTS_VALUE;
        pkt.dts = AV_NOPTS_VALUE;
        m_pYUVFrame->pts = videoStream->codec->frame_number;
        int ret = avcodec_encode_video2(ctx, &pkt, m_pYUVFrame, &got_output);
        if (ret < 0) {
            fprintf(stderr,"error encoding video frame: %s\n",av_err2str(ret));
        }
        if(got_output) {
            if(ctx->coded_frame->key_frame)
                pkt.flags|=AV_PKT_FLAG_KEY;
            pkt.stream_index = st->index;
            if(pkt.pts != AV_NOPTS_VALUE)
                pkt.pts = av_rescale_q(pkt.pts, videoStream->codec->time_base,videoStream->time_base);
            if(pkt.dts != AV_NOPTS_VALUE)
                pkt.dts = av_rescale_q(pkt.dts, videoStream->codec->time_base,videoStream->time_base);
            ret = av_interleaved_write_frame(fmtCtx, &pkt);
        } else {
            ret = 0;
        }
    }
    //clean
    av_frame_free(&m_pYUVFrame);
    clean();
}

void clean() {
    av_write_trailer(fmtCtx);
    /* Free the streams. */
    for (unsigned int i = 0; i< fmtCtx->nb_streams; i++) {
        //av_freep(&fmtCtx->streams->code);
        av_freep(&fmtCtx->streams);
    }
    if(!(fmtCtx->oformat->flags& AVFMT_NOFILE))
        /* Close the output file. */
        avio_close(fmtCtx->pb);
    /*free the stream */
    av_free(fmtCtx);
}

void print() {
    printf("this is codec file!");
}

#ifndef CODEC_TEST

int main() {
    print();
    work();
}

#endif // !CODEC_TEST
